#include <iostream>
#include <time.h>
#include <float.h>
#include <curand_kernel.h>
#include<tool.h>
#include<mytimer.h>
#include "sphere.cuh"
#include "hitable_list.cuh"
#include "camera.cuh"
#include"bvh.cuh"
#include"he.cuh"
#include"image_texture.cuh"
#include"rectangle.cuh"
#include"translate.cuh"
// limited version of checkCudaErrors from helper_cuda.h in CUDA examples



__device__ inline void write_color(unsigned char*fb,uchar3 color,int pixel_index)
{
    fb[pixel_index+0]=color.x;fb[pixel_index+1]=color.y;fb[pixel_index+2]=color.z;
}
__device__ vec3 color(const ray& r, vec3 background,bvh_node   *world, curandState *local_rand_state) {
    ray cur_ray = r;
    vec3 cur_attenuation = vec3(1.0,1.0,1.0);
    vec3 cur_emitted=vec3(0,0,0);
    float cur_pdf=1;
    
    for(int i = 0; i < 50; i++) 
    {
        
        hit_record rec;
        if (world->hit(cur_ray, 0.001f, FLT_MAX, rec)) 
        {
        //xz_rect(213, 343, 227, 332, 554




            ray scattered;
            vec3 attenuation;
            float pdf;
            vec3 emitted = rec.mat_ptr->emitted(rec.u, rec.v, rec.p);
            




            if(rec.mat_ptr->scatter(cur_ray, rec, attenuation, pdf,scattered, local_rand_state)) 
            {
                vec3 on_light = vec3(213 + curand_uniform(local_rand_state) * (343 - 213), 554, 227 + curand_uniform(local_rand_state) * (332 - 227));
                vec3 to_light = on_light - rec.p;
                float distance_squared = to_light.squared_length();
                to_light.make_unit_vector();
                if(dot(to_light,rec.normal)<0)
                {
                    return cur_emitted+cur_attenuation*emitted/cur_pdf;
                }
                float light_area = (343 - 213)*(332 - 227);
                float light_cosine = fabs(to_light.y());
                if (light_cosine < 1e-6)     
                {
                    return cur_emitted+cur_attenuation*emitted/cur_pdf;
                }
                pdf = distance_squared / (light_cosine * light_area);
                scattered=ray(rec.p,to_light,cur_ray.time());
                cur_pdf*=pdf;
                cur_attenuation *= attenuation;  
                //cur_attenuation*= rec.mat_ptr->scatter_pdf(cur_ray,rec,scattered);
                cur_ray = scattered;
            }
            else 
            {
                //return cur_emitted+cur_attenuation*emitted;
                return cur_emitted+cur_attenuation*emitted/cur_pdf;
                //return vec3(0.0,0.0,1.0);
            }
        }
        else 
        {//没射到返回背景值
            //vec3 c = vec3(0.5, 0.7, 1.0);
            vec3 c = background;
            //return vec3(1,1,1) * c;
            return cur_attenuation * c;
        }
    }
    
    return vec3(0.0,1.0,0.0); // exceeded recursion
}

__global__ void rand_init(curandState* rand_state) 
{
    if (threadIdx.x == 0 && blockIdx.x == 0) 
    {
        curand_init(1984, 0, 0, rand_state);
    }
}
__global__ void render_init(int max_x, int max_y, curandState *rand_state) 
{
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    int j = threadIdx.y + blockIdx.y * blockDim.y;
    if((i >= max_x) || (j >= max_y)) return;  
    int pixel_index = j*max_x + i;
    //Each thread gets same seed, a different sequence number, no offset
    curand_init(1984, pixel_index, 0, &rand_state[pixel_index]);
}
__global__ void test_curand(int max_x, int max_y, curandState *rand_state)
{
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    int j = threadIdx.y + blockIdx.y * blockDim.y;
    if ((i >= max_x) || (j >= max_y)) 
    {
       
        return;
    }
    
    int pixel_index = j*max_x + i;
    j=max_y-1-j;
   
    curandState local_rand_state = rand_state[pixel_index];
    float u = float(i + curand_uniform(&local_rand_state)) / float(max_x);
    float v = float(j + curand_uniform(&local_rand_state)) / float(max_y);

};
__global__ void test_render(unsigned char *fb, int max_x, int max_y, int ns, camera **cam, bvh_node  *world, curandState *rand_state)
{
    
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    int j = max_y-(threadIdx.y + blockIdx.y * blockDim.y)-1;
    if ((i >= max_x-1) || (j >= max_y-1)) 
    {
       
        return;
    }

    int pixel_index = ((max_y-j-1) * max_x + i);
    vec3 col(0,0,0);
    curandState local_rand_state = rand_state[0];
    float u = float(i + curand_uniform(&local_rand_state)) / float(max_x);
    float v = float(j + curand_uniform(&local_rand_state)) / float(max_y);
    ray r = (*cam)->get_ray(u,v,rand_state);
    ray cur_ray = r;
    
    
    color(cur_ray,vec3(0,0,0),world,rand_state);
    
    
    
}
__global__ void render(unsigned char *fb, int max_x, int max_y, int ns, vec3 background,camera **cam, bvh_node  *world, curandState *rand_state) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    int j = threadIdx.y + blockIdx.y * blockDim.y;
    if ((i >= max_x) || (j >= max_y)) 
    {
       
        return;
    }
    int pixel_index = j*max_x + i;
    j=max_y-1-j;
    
    curandState local_rand_state = rand_state[pixel_index];
    vec3 col(0,0,0);
    
    for(int s=0; s < ns; s++) 
    {
        float u = float(i + curand_uniform(&local_rand_state)) / float(max_x);
        float v = float(j + curand_uniform(&local_rand_state)) / float(max_y);
        ray r = (*cam)->get_ray(u,v,&local_rand_state);
        col += color(r, background,world, &local_rand_state);
    }
    
    
    
    col/=float(ns);
    
    col[0] = sqrt(col[0]);
    col[1] = sqrt(col[1]);
    col[2] = sqrt(col[2]);
    unsigned char c_r=static_cast<unsigned char>(clamp(col.x(),0,0.999999)*256);
    unsigned char c_g=static_cast<unsigned char>(clamp(col.y(),0,0.999999)*256);
    unsigned char c_b=static_cast<unsigned char>(clamp(col.z(),0,0.999999)*256);
    write_color(fb,make_uchar3(c_r,c_g,c_b),pixel_index*3);
}
const int sph_number=8 ;
const int image_width=800;
const int image_height=400;
const int sample_pix=100;
#define RND (curand_uniform(&local_rand_state))
__device__ void creat_CornellBox(hitable **d_list,camera **d_camera,int nx,int ny)
{
    auto red = new lambertian(new constant_texture(vec3(0.65, 0.05, 0.05)));
    auto white = new lambertian(new constant_texture(vec3(0.73, 0.73, 0.73)));
    auto green = new lambertian(new constant_texture(vec3(0.12, 0.45, 0.15)));
    auto light = new diffuse_light(new constant_texture(vec3(15, 15, 15)));
    int i=0;
    d_list[i++]=new yz_rect(0, 555, 0, 555, 555, green);//
    d_list[i++]=new yz_rect(0, 555, 0, 555, 0, red);
    d_list[i++]=new xz_rect(213, 343, 227, 332, 554, light);
    d_list[i++]=new xz_rect(0, 555, 0, 555, 555, white);
    d_list[i++]=new xz_rect(0, 555, 0, 555, 0, white);
    d_list[i++]=new xy_rect(0, 555, 0, 555, 555, white);
    
    const auto aspect_ratio = float(nx) / ny;
    vec3 lookfrom(278, 278, -800);
    vec3 lookat(278,278,0);
    vec3 vup(0,1,0);
    auto dist_to_focus = 10.0;
    auto aperture = 0.0;
    auto vfov = 40.0;
    *d_camera   = new camera(lookfrom,
                                 lookat,
                                 vup,
                                 aspect_ratio,
                                 vfov,
                                 aperture,
                                 dist_to_focus);

}
__global__ void create_world(hitable **d_list, 
    camera **d_camera, int nx, int ny, curandState *rand_state,aabb *aabb_bounding,cudaTextureObject_t texId) {
    if (threadIdx.x == 0 && blockIdx.x == 0) 
    {
        material* red = new lambertian(new constant_texture(vec3(0.65, 0.05, 0.05)));
        material* white1 = new lambertian(new constant_texture(vec3(0.73, 0.73, 0.73)));
        material* white2 = new lambertian(new constant_texture(vec3(0.73, 0.73, 0.73)));
        material* white3 = new lambertian(new constant_texture(vec3(0.73, 0.73, 0.73)));
        material* white4 = new lambertian(new constant_texture(vec3(0.73, 0.73, 0.73)));
        material* white5 = new lambertian(new constant_texture(vec3(0.73, 0.73, 0.73)));
        material* green = new lambertian(new constant_texture(vec3(0.12, 0.45, 0.15)));
        material* light = new diffuse_light(new constant_texture(vec3(15, 15, 15)));
        int i=0;
        d_list[i++]=new yz_rect(0, 555, 0, 555, 555, green);//1
        d_list[i++]=new yz_rect(0, 555, 0, 555, 0, red);//1
        d_list[i++]=new xz_rect(213, 343, 227, 332, 554, light);//1
        d_list[i++]=new xz_rect(0, 555, 0, 555, 555, white1);//1
        d_list[i++]=new xz_rect(0, 555, 0, 555, 0, white2);//
        d_list[i++]=new xy_rect(0, 555, 0, 555, 555, white3);//1
        hitable*box1=new box(vec3(0, 0, 0), vec3(165, 330, 165),white4);
        hitable*box2=new box(vec3(0,0,0), vec3(165,165,165),white5);
        d_list[i++]=new translate(new rotate_y(box1,15),vec3(265,0,295));
        d_list[i++]=new translate(new rotate_y(box2,-18),vec3(130,0,65));
        for(int i=0;i<sph_number;i++)
        {
            d_list[i]->bounding_box(0,1,aabb_bounding[i]);
            
        }
        const auto aspect_ratio = float(nx) / ny;
        vec3 lookfrom(278, 278, -800);
        vec3 lookat(278,278,0);
        vec3 vup(0,1,0);
        auto dist_to_focus = 10.0;
        auto aperture = 0.0;
        auto vfov = 40.0;
        *d_camera   = new camera(lookfrom,
                                    lookat,
                                    vup,
                                    aspect_ratio,
                                    vfov,
                                    aperture,
                                    dist_to_focus);
        

        
    }
}
__global__ void change_bvh(bvh_node*bvh)
{
    if (threadIdx.x == 0 && blockIdx.x == 0) 
    {
        bvh_node bb(*bvh);
        uint64_t funcaddr = *reinterpret_cast<uint64_t*>(&bb);
        memcpy(bvh,&funcaddr,sizeof(uint64_t));
    }
}
__global__ void free_world(hitable **d_list, camera **d_camera) 
{
    if (threadIdx.x == 0 && blockIdx.x == 0) 
    {
        for (int i = 0; i < sph_number; i++) 
        {
            delete d_list[i];
        }
        printf("delete d_list done");
        delete* d_camera;
    }
}
__global__ void test_bvh(bvh_node*bvh)
{
    printf("in kernel test_bvh\n");
    aabb temp;
    bvh->bounding_box(0,1,temp);
    vec3 a=temp.min();
    vec3 b=temp.max();
    printf("%f %f %f\n",a.x(),a.y(),a.z());
    printf("%f %f %f\n",b.x(),b.y(),b.z());
    printf("in kernel test_bvhend===========\n");
};

unsigned char* fb;
curandState *d_rand_state;
curandState* d_rand_state2;
hitable **d_list;
// sphere *d_sphere;
aabb*hitableBounds;
camera **d_camera;
bvh_node *d_bvh_node;
host_Manager*hostManager;
uint64_t sphhostaddr=0;
uint64_t hostaddr =0;
int tx = 8;
int ty = 8;
__host__ void init_cuda()
{
int num_pixels = image_width * image_height;
    size_t fb_size = num_pixels * sizeof(unsigned char)*3;
    hostManager=new host_Manager;

    // allocate FB
    checkCudaErrors(cudaMallocManaged((void **)&fb, fb_size));
    checkCudaErrors(cudaDeviceSynchronize());
    // allocate random state
    
    checkCudaErrors(cudaMalloc((void **)&d_rand_state, num_pixels*sizeof(curandState)));
    checkCudaErrors(cudaDeviceSynchronize());
    
    cudaMalloc((void**)&d_rand_state2, 1 * sizeof(curandState));
    checkCudaErrors(cudaDeviceSynchronize());
    //2nd random state is for the world
    rand_init << <1, 1 >> > (d_rand_state2);
    checkCudaErrors(cudaDeviceSynchronize());





    // make our world of hitables & the camera
    
    checkCudaErrors(cudaMallocManaged((void **)&d_list, sph_number*sizeof(hitable *)));
    checkCudaErrors(cudaDeviceSynchronize());
    
    checkCudaErrors(cudaMallocManaged((void **)&hitableBounds, sph_number*sizeof(aabb)));
    checkCudaErrors(cudaDeviceSynchronize());
    // sphhostaddr = *reinterpret_cast<uint64_t*>(&d_sphere[2]);
    
    
    checkCudaErrors(cudaMalloc((void **)&d_camera, sizeof(camera *)));
    checkCudaErrors(cudaDeviceSynchronize());
    create_world<<<1,1>>>(d_list, d_camera, image_width, image_height,d_rand_state2,hitableBounds,hostManager->get_texId());
    checkCudaErrors(cudaGetLastError());
    checkCudaErrors(cudaDeviceSynchronize());
    
    
    
    d_bvh_node=new bvh_node(d_list,sph_number,0,1,hitableBounds);
    
    hostaddr = *reinterpret_cast<uint64_t*>(d_bvh_node);
    checkCudaErrors(cudaGetLastError());
    checkCudaErrors(cudaDeviceSynchronize());
    change_bvh<<<1,1>>>(d_bvh_node);
    checkCudaErrors(cudaDeviceSynchronize());





}

__host__ void draw()
{
    dim3 blocks(image_width/tx+1,image_height/ty+1);
    dim3 threads(tx,ty);
    render_init<<<blocks, threads>>>(image_width, image_height, d_rand_state);
    checkCudaErrors(cudaGetLastError());
    checkCudaErrors(cudaDeviceSynchronize());
    vec3 background(0,0,0);
    
    render<<<blocks, threads>>>(fb, image_width, image_height,  sample_pix, background,d_camera, d_bvh_node, d_rand_state);
    checkCudaErrors(cudaGetLastError());
    checkCudaErrors(cudaDeviceSynchronize());
}
__host__ void clean_up()
{
    checkCudaErrors(cudaDeviceSynchronize());
    free_world<<<1,1>>>(d_list,d_camera);
    checkCudaErrors(cudaDeviceSynchronize());
    memcpy(d_bvh_node,&hostaddr,sizeof(uint64_t));
    delete d_bvh_node;
    checkCudaErrors(cudaGetLastError());
    checkCudaErrors(cudaFree(d_camera));
    // for(int i=0;i<sph_number;i++)
    // {
    //     memcpy(d_sphere+i,&sphhostaddr,sizeof(uint64_t));
    // }
    delete hostManager;
    checkCudaErrors(cudaFree(hitableBounds));
    
    checkCudaErrors(cudaFree(d_list));
    checkCudaErrors(cudaFree(d_rand_state));
    checkCudaErrors(cudaFree(d_rand_state2));
    checkCudaErrors(cudaFree(fb));
}
void run()
{
    std::cerr << "Rendering a " << image_width << "x" << image_height 
    << " image with " << sample_pix << " samples per pixel ";
    std::cerr << "in " << tx << "x" << ty << " blocks.\n";

    init_cuda();
    
    MyTimer mt;
    mt.Start();
    
    // Render our buffer
    draw();
    
    write("E:/workplace/c++/cudalearn/image.png",fb,image_width,image_height);
    mt.End();
    long gpuTime=mt.costTime;
    printf("Nvdia Execution and Trans time is: %ld ms \n",mt.costTime/1000);

    
    

    // clean up
    clean_up();
    

    cudaDeviceReset();
}
// int main() {
  
    
    

    
// }